import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import math
from torchvision.transforms.functional import InterpolationMode
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


def loadData():

    batch_size = 32
    num_workers = 2
    learning_rate=0.001
    epochs = 50  # 50轮
    # 判断是否有GPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 图像预处理
    mean = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
    std = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]

    transform_train = transforms.Compose([
        # transforms.RandomCrop(32, padding=4),
        # transforms.RandomHorizontalFlip(),
        # transforms.RandomRotation(15),
        # transforms.ToTensor(),
        # transforms.Normalize(mean, std)
        transforms.Resize((128, 128)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    transform_test = transforms.Compose(
        [
         transforms.Resize((128, 128)),
         transforms.ToTensor(),
         transforms.Normalize(mean, std)])


    # CIFAR-100 数据集下载
    train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                                 train=True,
                                                 transform=transform_train,
                                                 download=True)

    test_dataset = torchvision.datasets.CIFAR100(root='data/',
                                                train=False,
                                                transform=transform_test)

    # 数据载入
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               num_workers=num_workers,
                                               shuffle=True)

    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              num_workers=num_workers,
                                              shuffle=False)

    for batch_idx, (data, target) in enumerate(train_loader):
        print(data.shape)
        print(target.shape)
        break

    return train_loader,test_loader

# if __name__ == '__main__':
#     loadData()
